﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace CEED.Utils
{
    public enum NcchSection
    {
        exheader = 1,
        exefs,
        romfs
    }

    public class NCCHInfo
    {
        List<NCCHInfoEntry> Entries = new List<NCCHInfoEntry>();
        public SEEDinfo seedinfo;
        public static SEEDEntry SEED;
        public NCCHInfo(NCCH ncch, int PartitionIndex, Encoding encoding)
        {
            Entries = new List<NCCHInfoEntry>();
            int Use7x = ncch.Header.Flags[3];

            if ((ncch.Header.Flags[7] & 0x20) == 0x20)
            {
                if (!File.Exists("seedinfo.bin"))
                    throw new Exception("Cannot find seedinfo.bin.");
                seedinfo = new SEEDinfo("seedinfo.bin");
                for (int i = 0; i < seedinfo.Count; i++)
                {
                    if (seedinfo.Entries[i].Checksum == ncch.Header.SEEDCheck)
                    {
                        byte[] KeyY = new byte[16];
                        Array.Copy(ncch.Header.NCCHHeaderSignature, 0, KeyY, 0, 16);
                        seedinfo.Entries[i] = new SEEDEntry(seedinfo.Entries[i], KeyY);
                        SEED = seedinfo.Entries[i];
                        break;
                    }
                }
            }


            if (ncch.Header.ExtendedHeaderSize > 0)
            {
                Entries.Add(new NCCHInfoEntry(ncch, PartitionIndex, NcchSection.exheader, encoding, 0, "exheader"));
            }
            if (ncch.Header.ExeFSLength > 0)
            {
                Entries.Add(new NCCHInfoEntry(ncch, PartitionIndex, NcchSection.exefs, encoding, 0, "exefs_norm"));
                if (Use7x == 1)
                {
                    Entries.Add(new NCCHInfoEntry(ncch, PartitionIndex, NcchSection.exefs, encoding, Use7x, "exefs_7x"));
                }
            }
            if (ncch.Header.RomFSLength > 0)
            {
                Entries.Add(new NCCHInfoEntry(ncch, PartitionIndex, NcchSection.romfs, encoding, Use7x, "romfs"));
            }
        }
        public void ToFile(string path)
        {
            using (FileStream fs = File.Create(path))
            {
                BinaryWriter writer = new BinaryWriter(fs);
                writer.Write(0xFFFFFFFF);
                writer.Write(0xF0000003);
                writer.Write(Entries.Count);
                writer.Write((uint)0);
                foreach (NCCHInfoEntry nie in Entries)
                {
                    nie.WriteToStream(fs);
                }
            }
        }

        public static void GetNcchAesCounter(ref byte[] counter, NCCH ncch, NcchSection type)
        {
            counter = new byte[16];
            if (ncch.Header.Version == 2 || ncch.Header.Version == 0)
            {
                Array.Copy(ReverseByteArray(ncch.Header.TitleID), 0, counter, 0, 8);
                counter[8] = (byte)type;
            }
            else if (ncch.Header.Version == 1)
            {
                long num = 0;
                switch (type)
                {
                    case NcchSection.exheader:
                        num = 0x200;
                        break;
                    case NcchSection.exefs:
                        num = ncch.Header.ExeFSOffset * 0x200;
                        break;
                    case NcchSection.romfs:
                        num = ncch.Header.RomFSOffset * 0x200;
                        break;
                }
                Array.Copy(ReverseByteArray(ncch.Header.TitleID), 0, counter, 0, 8);
                for (int i = 0; i < 4; i++)
                {
                    counter[i + 12] = (byte)((num >> ((3 - i) * 8)) & 0xFF);
                }
            }
        }
        public static byte[] ReverseByteArray(byte[] source)
        {
            return source.Reverse<byte>().ToArray<byte>();
        }
        public static int roundUp(int numToRound, int multiple)
        {
            if (multiple == 0)
                return numToRound;

            int remainder = Math.Abs(numToRound) % multiple;
            if (remainder == 0)
                return numToRound;
            if (numToRound < 0)
                return -(Math.Abs(numToRound) - remainder);
            return numToRound + multiple - remainder;
        }
    }
    public class NCCHInfoEntry
    {
        public readonly string[] PartitionName = new string[] { "Main", "Manual", "DownloadPlay", "Partition4", "Partition5", "Partition6", "Partition7", "UpdateData" };
        public byte[] Counter = new byte[16];
        public byte[] KeyY = new byte[16];
        public int Size;
        public byte[] Reserved = new byte[8];
        public int Use7x = 0;
        public byte[] FileName = new byte[112];

        public NCCHInfoEntry(NCCH ncch, int PartitionIndex, NcchSection type, Encoding encoding, int use7x, string sectionname)
        {
            NCCHInfo.GetNcchAesCounter(ref Counter, ncch, type);
            Array.Copy(ncch.Header.NCCHHeaderSignature, 0, KeyY, 0, 16);
            if ((ncch.Header.Flags[7] & 0x20) == 0x20)
            {
                if ((type == NcchSection.romfs) || ((type == NcchSection.exefs) && (use7x == 1)))
                {
                    Array.Copy(NCCHInfo.SEED.KeyY, 0, KeyY, 0, 16);
                }
            }

            switch (type)
            {
                case NcchSection.exefs:
                    Size = NCCHInfo.roundUp((int)ncch.Header.ExeFSLength * 0x200, 1024 * 1024) / (1024 * 1024);
                    break;
                case NcchSection.exheader:
                    Size = NCCHInfo.roundUp((int)ncch.Header.ExtendedHeaderSize * 0x200, 1024 * 1024) / (1024 * 1024);
                    break;
                case NcchSection.romfs:
                    Size = NCCHInfo.roundUp((int)ncch.Header.RomFSLength * 0x200, 1024 * 1024) / (1024 * 1024);
                    break;
            }
            using (MemoryStream ms = new MemoryStream())
            {
                BinaryWriter writer = new BinaryWriter(ms, encoding);
                string path = "";
                if (encoding == Encoding.Unicode)
                    path = string.Format("sdmc:/{0:X16}.{1}.{2}.xorpad", BitConverter.ToUInt64(ncch.Header.TitleID, 0), PartitionName[PartitionIndex], sectionname);
                else if (encoding == Encoding.UTF8)
                    path = string.Format("/{0:X16}.{1}.{2}.xorpad", BitConverter.ToUInt64(ncch.Header.TitleID, 0), PartitionName[PartitionIndex], sectionname);
                writer.Write(path.ToCharArray());
                Array.Copy(ms.ToArray(), 0, FileName, 0, ms.Length);
            }

            Use7x = use7x;
        }
        public void WriteToStream(Stream stream)
        {
            BinaryWriter writer = new BinaryWriter(stream);
            writer.Write(Counter);
            writer.Write(KeyY);
            writer.Write(Size);
            writer.Write(Reserved);
            writer.Write(Use7x);
            writer.Write(FileName);
        }
    }
}
